import torch
from torch import Tensor, randint
from torchmetrics import Metric
from torchmetrics.functional.segmentation import hausdorff_distance


class HD95Metric(Metric):
    def __init__(self, num_classes: int = 1):
        super().__init__()
        self.num_classes = num_classes
        self.add_state("distance", default=torch.tensor(0.0), dist_reduce_fx="sum")
        self.add_state("total", default=torch.tensor(0), dist_reduce_fx="sum")

    def update(self, pred: Tensor, target: Tensor) -> None:
        if pred.shape != target.shape:
            raise ValueError("preds and target must have the same shape")
        pred = pred.unsqueeze(1) > 0.5
        target = target.unsqueeze(1)

        data = hausdorff_distance(pred, target, num_classes=self.num_classes) * 0.95

        self.distance += data[0][0]
        self.total += 1

    def compute(self) -> Tensor:
        return self.distance / self.total
